package org.apache.lucene.ocean;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;

import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.SerialMergeScheduler;
import org.apache.lucene.store.Directory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Allows creation of an index using multiple threads by feeding documents into
 * a BlockingQueue.  This is to make it easier to 
 * 
 */
// TODO: after create called make object unusable
public class IndexCreator {
  final static Logger LOG = LoggerFactory.getLogger(IndexCreator.class);
  private Long maxSize;
  private int threads;
  private ExecutorService threadPool;
  private IndexWriter indexWriter;
  private boolean isFinished;
  private List<Future<Object>> futures = new ArrayList<Future<Object>>();
  private BlockingQueue<Add> queue;
  private int numAdded = 0;
  private String name;
  
  /**
   * 
   * @param name Name used for debugging
   * @param directory
   * @param maxSize
   * @param threads Number of threads to use
   * @param analyzer Default analyzer
   * @param threadPool 
   * @throws IOException
   */
  public IndexCreator(String name, Directory directory, long maxSize, int threads, Analyzer analyzer, ExecutorService threadPool) throws IOException {
    this.name = name;
    this.maxSize = maxSize;
    this.threads = threads;
    this.threadPool = threadPool;
    isFinished = false;
    indexWriter = new IndexWriter(directory, false, analyzer, true);
    indexWriter.setUseCompoundFile(false);
    indexWriter.setMergeScheduler(new SerialMergeScheduler());
    indexWriter.setMaxBufferedDocs(Integer.MAX_VALUE);
    indexWriter.setRAMBufferSizeMB(IndexWriter.DISABLE_AUTO_FLUSH);
  }
  
  public IndexCreator(IndexWriter indexWriter, int threads, ExecutorService threadPool) throws IOException {
    this.indexWriter = indexWriter;
    this.threads = threads;
    this.threadPool = threadPool;
    isFinished = false;
  }
  
  /**
   * Close the indxWriter, setFinished to true
   * @throws Exception
   */
  public void close() throws Exception {
    indexWriter.close();
    setFinished(true);
  }
  
  /**
   * Number of documents added
   */
  public int getNumAdded() {
    return numAdded;
  }
  
  /**
   * Add document object.  Wraps a document and an analyzer.
   *
   */
  public static class Add {
    private final Document document;
    private final Analyzer analyzer;

    public Add(Document document) {
      this.document = document;
      this.analyzer = null;
    }
    
    public Add(Document document, Analyzer analyzer) {
      this.document = document;
      this.analyzer = analyzer;
    }
    
    public Analyzer getAnalyzer() {
      return analyzer;
    }
    
    public Document getDocument() {
      return document;
    }
  }
  
  /**
   * Starts the tasks in threads
   * @param queue
   * @throws Exception
   */
  public void start(BlockingQueue<Add> queue) throws Exception {
    this.queue = queue;
    
     // set impossibly high to never be
    // triggered, setting both to
    // DISABLE_AUTO_FLUSH causes an
    // exception
    // List<Callable<Object>> callables = new
    // ArrayList<Callable<Object>>(threads);
    for (int x = 0; x < threads; x++) {
      // callables.add(new Task(queue));
      futures.add(threadPool.submit(new Task()));
    }
    // futures = threadPool.invokeAll(callables);

  }
  
  /**
   * Create the index
   * @param optimize If true, optimize the index at completion
   * @throws Exception
   */
  public void create(boolean optimize) throws Exception {
    while (queue.peek() != null) {
      Thread.sleep(5);
    }
    setFinished(true);
    try {
      for (Future<Object> future : futures) {
        if (future.isDone()) {
          try {
            future.get();
          } catch (ExecutionException executionException) {
            Throwable cause = executionException.getCause();
            if (cause instanceof Exception) 
              throw (Exception) cause;
            else 
              throw new Exception(cause);
          }
        }
        Thread.sleep(10);
      }
      if (optimize) indexWriter.optimize(); // should not be necessary
      indexWriter.commit();
      numAdded = indexWriter.numDocs();
    } finally {
      indexWriter.close();
    }
  }

  public void setFinished(boolean isFinished) {
    this.isFinished = isFinished;
  }
  
  /**
   * If the index creation is finished.
   * @return
   */
  private boolean isFinished() {
    if (isFinished)
      return true;
    if (indexWriter.ramSizeInBytes() >= maxSize) {
      isFinished = true;
    }
    return isFinished;
  }
  
  /**
   * Task that executes in a thread.  Polls the queue for new Add objects
   * and calls addDocument on the indexWriter.
   *
   */
  public class Task implements Callable {
    public Object call() throws Exception {
      while (!isFinished()) {
        Add add = queue.poll();
        if (add != null) {
          Document document = add.getDocument();
          Analyzer analyzer = add.getAnalyzer();
          if (analyzer == null) {
            indexWriter.addDocument(document);
          } else {
            indexWriter.addDocument(document, analyzer);
          }
        } else {
          Thread.sleep(5);
        }
      }
      return null;
    }
  }
}
